import os
import sys
import numpy as np
import matplotlib.pyplot as plt
import IONData
import SwcLoader
import Scene
import GeometryAdapter
import Visual as nv
from distinctipy import distinctipy
import nrrd
import BoundLaplace as BL
import Flatmap
from collections import Counter
import NeuronProcess
import BrainRegion
from tqdm.notebook import tqdm as tool_bar
from sklearn.cluster import AgglomerativeClustering
import seaborn
from PIL import Image
import cv2 as cv
import copy
from scipy.cluster import hierarchy
from sklearn.manifold import TSNE
import matplotlib as mpl
import matplotlib.ticker as ticker
from cycler import cycler
import math
from sklearn.cluster import DBSCAN
import networkx as nx
from netgraph import Graph
from scipy import ndimage
from scipy import stats
import mplcyberpunk
from matplotlib.colors import LinearSegmentedColormap
#import tensorflow as tf
import matplotlib.colors as mcolors
import matplotlib.cm as cm
import matplotlib.patches as mpatches
from scipy import stats

'''Variables'''
fontsize_text = 20*0.7
fontsize_ticks = 20*0.7
fontsize_label = 25*0.7
fontsize_title = 30*0.7
#plt.style.use('cyberpunk')
console_channel_1 = sys.stdout
console_channel_2 = open(os.devnull,'w')
iondata = IONData.IONData()
br = BrainRegion.BrainRegion()
br.praseJson()
background_colors = np.array([[0.0,0.0,0.0],[1.0,1.0,1.0]])
calculating_depth_flatten_position_mode = 1 # 0-> calculating 1-> look table
if calculating_depth_flatten_position_mode==0:
    res,gridpath = iondata.getFileFromServer("boundlaplace20.nrrd")
    grid,header = nrrd.read(gridpath)
    resRelaxation,RelaxationPath=iondata.getFileFromServer('boundlaplaceout20.nrrd')
    relaxation,relaxationheader = nrrd.read(RelaxationPath)
    resdv0, dv0Path = iondata.getFileFromServer('dv0.nrrd')
    dv0, dv0header = nrrd.read(dv0Path)
    resdv1, dv1Path = iondata.getFileFromServer('dv1.nrrd')
    dv1, dv1header = nrrd.read(dv1Path)
    resdv2, dv2Path = iondata.getFileFromServer('dv2.nrrd')
    dv2, dv2header = nrrd.read(dv2Path)
    dv0 = dv0.astype(np.float32) / 1000 - 1
    dv1 = dv1.astype(np.float32) / 1000 - 1
    dv2 = dv2.astype(np.float32) / 1000 - 1 
    flatenPara = Flatmap.createSurfaceGraph()
else:
    depth20,header = nrrd.read('../../resource/depth20.nrrd')
    outpx,header = nrrd.read('../../resource/outpx.nrrd')
    outpy,header = nrrd.read('../../resource/outpy.nrrd')
    outp1,header = nrrd.read('../../resource/outP1.nrrd')
    outp2,header = nrrd.read('../../resource/outP2.nrrd')
    outp3,header = nrrd.read('../../resource/outP3.nrrd')
    inp1,header = nrrd.read('../../resource/inP1.nrrd')
    inp2,header = nrrd.read('../../resource/inP2.nrrd')
    inp3,header = nrrd.read('../../resource/inP3.nrrd')
def rgb_to_hex(rgb):
    hex_value = '#{:02x}{:02x}{:02x}'.format(int(rgb[0]),int(rgb[1]),int(rgb[2]))
    return hex_value
color_pool = np.array([
        [230, 25, 75],     # 红色
        [60, 180, 75],     # 绿色
        [255, 225, 25],    # 黄色
        [0, 130, 200],     # 蓝色
        [245, 130, 48],    # 橙色
        [145, 30, 180],    # 紫色
        [70, 240, 240],    # 青色
        [240, 50, 230],    # 粉色
        [210, 245, 60],    # 淡黄色
        [250, 190, 190],   # 浅红色
        [0, 128, 128],     # 青绿色
        [230, 190, 255],   # 淡紫色
        [170, 110, 40],     # 棕色
        [192, 192, 192],  # 银色 
        [152, 138, 255],
        [175, 255, 131],
        [8, 0, 255],
        [212, 255, 178]
    ])/255
    #colors = np.random.choice(colors,replace=False)
color_pool_hex = []
for i in color_pool:
    color_pool_hex.append(rgb_to_hex(np.round(i*255).astype(int)))
hierarchy.set_link_color_palette(color_pool_hex)
cortex_region = ['FRP','PL','ILA','ORB','ACA','AI','GU','VISC','PTLp','RSP','TEa','ECT','PERI','MOs','MOp','SSp','SSs','AUD','VIS']
brain_hierachy = {
    'Cortex':['FRP','PL','ILA','ORB','ACA','AI','GU','VISC','PTLp','RSP','TEa','ECT','PERI','MOs','MOp','SSp','SSs','AUD','VIS'],
    'OLF':['MOB','AOB','AON','TT','DP','PIR','NLOT','COA','PAA','TR'],
    'HPF':['CA','DG','FC','IG','ENT','PAR','POST','PRE','SUB','ProS','HATA','APr'],
    'Cortex-sp':['CLA','EP','LA','BLA','BMA','PA'],
    'CNU':['STRd','ACB','FS','OT','LSX','AAA','CEA','MEA','PALd','PALv','PALm','PALc'],
    'TH':['VAL','VM','VPL','VPM','VPLpc','VPMpc','PoT','SPA','SPFm','SPFp','PP','MG','LGd','LP','PO','POL','SGN','Eth','AV','AM','AD','IAM','IAD','LD',
    'IMD','MD','SMT','PR','PVT','PT','RE','Xi','RH','CM','PCN','CL','PF','PIL','GENv','EPI','RT'],
    'HY':['SO','ASO','PVH','PVa','PVi','ARH','ADP','AVP','AVPV','DMH','MEPO','MPO','OV','PD','PS','PVp','PVpo','SBPV','SCH','SFO','VMPO','VLPO','AHN',
    'MBO','MPN','PMd','PMv','PVHd','VMH','PH','LHA','LPO','PST','PSTN','PeF','RCH','STN','TU','ZI','ME'],
    'MB':['APN','NOT','NPC','PPT','RPF','AT','MPT','OP','VTA','SNr','SNc','SCs','SCm','NB','SAG','PBG','IC','RR','MRN','RN',
    'PAG','PRC','INC','ND','Su3','CUN','PPN','IF','IPN','RL','CLI','DR','III','MA3','EW','VTN'],
    'P':['PSV','PB','PCG','PG','PRNr','PRNc','SUT','TRN','V','P5','CS','LDT','NI','RPO','SOC','NLL','DTN','PDTg','SG','I5','SLC','SLD'],
    'MY':['DCO','VCO','LAV','MV','SPIV','SUV','DMX','GR','CU','ECU','NTS','NTB','SPVC','SPVI','SPVO','Pa5','VII','GRN','IRN','LIN','LRNm','LRNp','MARN',
    'MDRNd','MDRNv','PARN','PGRNd','PGRNl','PHY','VI','AMB','ICB','PAS','PRP','PPY','x','XII','y','IO','RM','RPA','RO'],
    'CB':['VERM','HEM','FN','IP','DN','VeCB'],
    'FT':['py','pyd'],
    'SC':['unknow']
}
brain_region = []
for a in brain_hierachy:
    for b in brain_hierachy[a]:
        brain_region.append(b)
def get_subregions(regions):
    a = {}
    b = {}
    for i in regions:
        a[i] = []
        for j in br.getRegionList(i):
            if j[1]!=0:
                a[i].append(j[0])
                b[j[0]] = [i][0]
    return a,b
cortex_region_to_subregion,cortex_subregion_to_region = get_subregions(cortex_region)
_wedgeprops = {"edgecolor" : "black", 'linewidth': 1,'antialiased': True}
_flatmap = cv.imread('../../resource/flatmapedge2.tif',1)
_flatmap_mask = cv.resize(cv.imread('../../resource/flatmap.png',1),(_flatmap.shape[0],_flatmap.shape[1]))
_flatmap_name = cv.resize(cv.imread('../../resource/flatmapedge.png',1),(_flatmap.shape[0],_flatmap.shape[1]))
flatmap_mask = np.zeros((_flatmap.shape[0],_flatmap.shape[1],4))
flatmap_mask[...,:3] = (_flatmap_mask==0)*255
flatmap_mask[...,3] = np.where(flatmap_mask[...,0]==0,0,1)
flatmap = np.zeros((_flatmap.shape[0],_flatmap.shape[1],4))
flatmap[...,-1] = np.where(_flatmap_mask[...,0]==0,1.0,flatmap[...,-1])
flatmap[...,0] = 255-(np.mean(_flatmap_name,axis=-1)<200)*255
flatmap[...,1] = 255-(np.mean(_flatmap_name,axis=-1)<200)*255
flatmap[...,2] = 255-(np.mean(_flatmap_name,axis=-1)<200)*255
flatmap[...,3] = np.mean(_flatmap_name,axis=-1)<200
flatmap2 = np.zeros((_flatmap.shape[0],_flatmap.shape[1],4))
flatmap2[...,-1] = np.where(_flatmap_mask[...,0]==0,1.0,flatmap[...,-1])
flatmap2[...,0] = 255-(np.mean(_flatmap,axis=-1)<200)*255
flatmap2[...,1] = 255-(np.mean(_flatmap,axis=-1)<200)*255
flatmap2[...,2] = 255-(np.mean(_flatmap,axis=-1)<200)*255
flatmap2[...,3] = np.mean(_flatmap,axis=-1)<200
'''System control'''
def control_console_output(state):
    if state:
        sys.stdout = console_channel_1
    else:
        sys.stdout = console_channel_2
'''math'''
def func_sigmoid(x):
    return 1 / (1 + np.exp(-x))
def func_gaussian(x,center=0,sigma=1):
    return (1/(2*np.pi*sigma**2))*np.exp(-((x-center)**2)/(2*sigma**2))
'''Data processing'''
def find_top_n_indices(arr,n):
    sorted_indices = sorted(enumerate(arr),key=lambda x: x[1],reverse=True)
    top_n_indices = [index for index,value in sorted_indices[:n]]
    return top_n_indices
def show_cmap(fig,cmap,norm=None,extend=None):
    if norm is None:
        norm = mcolors.Normalize(vmin=0, vmax=cmap.N)
    im = cm.ScalarMappable(norm=norm, cmap=cmap)
    fig.colorbar(im, cax=ax,orientation='vertical', extend=extend)
    plt.show()
def plot_transition_color(transition_color,X):
    transition_color_step = 1.0/(len(transition_color)-1)
    X = np.expand_dims(X/np.max(X),axis=-1)
    last_mask = 0
    result = []
    for i in range(1,len(transition_color)):
        mask = (X<transition_color_step*i)-last_mask
        result.append(mask*(transition_color[i]*(X-transition_color_step*(i-1))/transition_color_step+transition_color[i-1]*(transition_color_step*i-X)/transition_color_step))
        last_mask += mask
    return np.sum(result,axis=0)
'''Connection density estimate'''
def generate_gaussian_kernel(r,s):
    x = np.linspace(-r, r, 2*r+1)
    y = np.linspace(-r, r, 2*r+1)
    x, y = np.meshgrid(x, y)
    return np.exp(-(x**2 + y**2) / (2*s**2))
def fill_kernel(X,p,c,r,kernel):
    tmp = np.array([np.round(p).astype(int)-r,np.round(p).astype(int)+r+1])
    p = np.array([
        np.where(tmp[0]<0,0,tmp[0]),
        np.array([
            np.min([tmp[1][0],X.shape[0]+1]),
            np.min([tmp[1][1],X.shape[1]+1])
            ])
        ])
    s = p-tmp
    s[1][0] = s[1][0] if s[1][0]<0 else kernel.shape[0]
    s[1][1] = s[1][1] if s[1][1]<0 else kernel.shape[1]
    X[p[0,0]:p[1,0],p[0,1]:p[1,1]] += c*kernel[s[0,0]:s[1,0],s[0,1]:s[1,1]]
def plot_connection_density(info,r = 200,s = 125,transition_color = np.array([[255,255,255],[161,238,204],[255,242,193],[246,49,141]])/255.0,save_path='test.pdf'):
    kernel = generate_gaussian_kernel(r,s)
    cmap1 = mcolors.LinearSegmentedColormap.from_list('cmap1', transition_color)
    fig = plt.figure(figsize=(5,5))
    tmp = _flatmap[...,0]*0.0
    soma_list = [[],[]]
    for k in info:
        p = info[k]['soma_flatten_position']
        e = info[k]['projection_strength']
        fill_kernel(tmp,np.array(p[::-1])*2,e/info[k]['relative_density'],r,kernel)
        soma_list[0].append(np.array(p[::-1])*2)
        soma_list[1].append(e)
    soma_list[1] = np.array(soma_list[1])**0.5
    soma_list[1] /= np.max(soma_list[1])
    result = np.where(_flatmap[...,0]==-1,0,tmp)
    for k,l in zip(soma_list[0],soma_list[1]):
        cv.circle(result,[int(np.round(k[1])), int(np.round(k[0]))],int(np.round(l*10))+1,np.array([0.5,0.5,0.5]),-1)
    result /= np.max(result)
    _ = plt.imshow(result,cmap=cmap1,vmin=0,vmax=1)
    plt.imshow(flatmap_mask,zorder=1)
    plt.imshow(flatmap2,zorder=2)
    plt.xlim(flatten_map_r[0],flatten_map_r[1])
    plt.ylim(flatten_map_r[2],flatten_map_r[3])
    plt.axis('off')
    fig.colorbar(_,extend='both')
    plt.suptitle('Connection density')
    plt.rcParams['pdf.fonttype'] = 42
    plt.rcParams['ps.fonttype'] = 42
    plt.savefig(save_path,format='pdf',bbox_inches='tight',dpi=500)
    plt.show()
'''Others'''
def get_contour(X):
    X = np.array(X).astype(float)
    filter = np.ones((3,3,1,1))
    tmp = tf.nn.conv2d(X.reshape(1,X.shape[0],X.shape[1],1),filter,(1,1,1,1),padding='SAME').numpy().reshape(X.shape[0],X.shape[1])
    return np.where(tmp==X*9,0.0,1.0)
def plot_points_distribution_in_mask(X,mask,title='',save_path='tmp',cmap='Reds'):
    a = seaborn.kdeplot(X[:,1],X[:,0],fill=True,cmap=cmap,common_norm=True,cumulative=False)
    b = np.array([1-mask,1-mask,1-mask,1-mask]).transpose(1,2,0)
    a.imshow(b,zorder=1)
    plt.title(title,color='k')
    plt.axis('off')
    plt.axis('off')
    plt.axis("equal")
    plt.xlim(0,mask.shape[1])
    plt.ylim(mask.shape[0],0)
    plt.savefig(save_path+'.pdf',format='pdf')
    plt.savefig(save_path+'.png',format='png')
    plt.show()
def plot_post_process(title='',xlabel='',ylabel='',save_path=''):
    plt.title(title,fontsize=fontsize_title,color='k')
    plt.xlabel(xlabel,fontsize=fontsize_label,color='k')
    plt.ylabel(ylabel,fontsize=fontsize_label,rotation=-90,color='k').set_verticalalignment('top')
    plt.xticks(rotation = -90)
def normalization(x):
    y = np.array(x)
    y_norm = np.linalg.norm(y,axis=-1).reshape(-1,1)
    return [y_norm,y/np.where(y_norm==0.0,1.0,y_norm)]
def pie_chart_label_func(pct,allvals):
    absolute = int(np.round(pct/100.*np.sum(allvals)))
    return f"{pct:.1f}%\n({absolute:d})"
def estimate_density_function(X,parameter=[0,1200,1,35],weight=[]): #start,length,step,kernel_sigma
    distribution_length = parameter[1]
    distribution_discrete = np.zeros((parameter[1],))
    distribution_continuous = np.zeros((parameter[1],))
    distribution_index = np.array([i*parameter[2]+parameter[0] for i in range(distribution_length)])
    for index,i in enumerate(X):
        center_continuous = np.clip(i,parameter[0],parameter[0]+parameter[1]*parameter[2])
        center_discrete = np.clip(int(np.round((center_continuous-parameter[0])/parameter[2])),0,parameter[1]-1)
        g = func_gaussian(distribution_index,center_continuous,sigma=parameter[3])
        distribution_continuous += g/np.sum(g)*(1.0 if len(weight)==0 else weight[index]) 
        distribution_discrete[center_discrete] += 1*(1.0 if len(weight)==0 else weight[index]) 
    distribution_continuous /= np.sum(distribution_continuous)
    distribution_discrete /= np.sum(distribution_discrete)
    return [distribution_discrete,distribution_continuous]
def subcube(x,center,r):
    slc = tuple(slice(max(c - r, 0), c + r + 1) for c in center)
    subcube = x[slc]
    padded_subcube = np.pad(subcube,[(max(r-center[i],0),max(center[i]+r-j+1,0)) for i,j in enumerate(x.shape)],mode='constant')
    return padded_subcube
def create_dir(directory):
    if not os.path.exists(directory):
        os.mkdir(directory)
    else:
        pass
def get_neuron_tree(neuron):
    return iondata.getNeuronTreeByID(neuron[:-7],neuron[-7:])
def get_region_volume(regions):
    _br = copy.deepcopy(br)
    result = []
    for i in regions:
        tmp = _br.getRegion(i)
        tmp.annotation()
        result.append(np.sum(tmp.readdata))
        del tmp.readdata
    return result
def get_region_volume_and_center_and_range(regions):
    _br = copy.deepcopy(br)
    result = [[],[],[]]
    for i in regions:
        tmp = _br.getRegion(i)
        tmp.annotation()
        result[0].append(np.sum(tmp.readdata))
        tmp2 = np.transpose(np.nonzero(tmp.readdata[...,:tmp.readdata.shape[2]//2]))
        result[1].append(np.mean(tmp2, axis=0))
        position_range_info = [np.min(tmp2,axis=0),np.max(tmp2,axis=0)]
        result[2].append(position_range_info)
        del tmp.readdata
    return result
def get_mirrored_point(p,mode=0,control=0):
    y = np.array(p)
    if mode==0:
        if y[2] > 5700:
            y[2] = 11400-y[2]
        return y
    elif mode==1:
        if control:
            y[2] = 11400-y[2]
        return y
    else:
        return y
def draw_line_in_mask(mask,terminal,mirror=False,scale=10):
    x = terminal
    while x!=None:
        p = (np.array(x.xyz)/scale).astype(int)
        if mirror:
            mask[p[0],p[1],1140-p[2]] += 1
        else:
            mask[p[0],p[1],p[2]] += 1
        x = x.parent
def get_mask_value(mask,point):
    tmp = (np.round(point)).astype(int)
    tmp = np.where(tmp<0,0,tmp)
    tmp2 = tmp-np.array(mask.shape)+1
    tmp2 = np.where(tmp2>0,tmp2,0)
    tmp = tmp-tmp2
    return mask[tmp[0],tmp[1],tmp[2]]
def if_inside_mask(mask,point,r1=0,r2=1):
    try:
        tmp = (np.round(point)).astype(int)
        tmp = np.array([tmp-r1,tmp+r2])
        tmp = np.where(tmp<0,0,tmp)
        return 1 if np.any(mask[tmp[0,0]:tmp[1,0],tmp[0,1]:tmp[1,1],tmp[0,2]:tmp[1,2]]>0) else 0
    except:
        return 0
def position_color_mapping(x,dimension=3,position_colors = np.array([[[1,0,0],[0,1,1]],[[0,1,0],[1,0,1]],[[0,0,1],[1,1,0]]])):
    position_range_info = [np.min(x,axis=0),np.max(x,axis=0)]
    position_ratio = (x-position_range_info[0])/(position_range_info[1]-position_range_info[0])
    result = 0
    for i in range(dimension):
        result += (position_colors[i,0:1]*position_ratio[:,i:i+1]+position_colors[i,1:2]*(1-position_ratio[:,i:i+1]))/3
    for i in range(3):
        result[:,i] = func_sigmoid((result[:,i]-np.mean(result[:,i]))/np.std(result[:,i]))
    return result
def get_value_range(X):
    try:
        tmp = (np.sum(X,axis=0)>0).tolist()
        x_r = [tmp.index(True),len(tmp)-tmp[::-1].index(True)]
    except:
        x_r = [-1,-1]
    try:
        tmp = (np.sum(X,axis=1)>0).tolist()
        y_r = [tmp.index(True),len(tmp)-tmp[::-1].index(True)]
    except:
        y_r = [-1,-1]
    return [x_r,y_r]
a,b = get_value_range(1-flatmap[...,0]/255)
tmp = [np.mean(a),np.mean(b)]
tmp2 = np.max([a[1]-a[0],b[1]-b[0]])
flatten_map_r = [tmp[0]-tmp2/2,tmp[0]+tmp2/2,tmp[1]+tmp2/2,tmp[1]-tmp2/2]
def plot_points_distribution_in_mask(X,mask,title='',save_path='tmp',cmap='Reds'):
    a = seaborn.kdeplot(X[:,1],X[:,0],fill=True,cmap=cmap,common_norm=True,cumulative=False)
    b = np.array([1-mask,1-mask,1-mask,1-mask]).transpose(1,2,0)
    a.imshow(b,zorder=1)
    plt.title(title,color='k')
    plt.axis('off')
    plt.axis('off')
    plt.axis("equal")
    a,b = get_value_range(mask)
    plt.xlim(a[0],a[1])
    plt.ylim(b[1],b[0])
    plt.savefig(save_path+'.pdf',format='pdf')
    plt.savefig(save_path+'.png',format='png')
    plt.show()
def get_axonal_statics_by_regions(neuron,regions,search_key='projectregion'):
    neuron_property = iondata.getNeuronPropertyByID(neuron[:-7],neuron[-7:])
    brproperty = BrainRegion.RegionProperty(copy.deepcopy(br))
    brproperty.setProperty(neuron_property[search_key])
    tmp = []
    for region in regions:
        tmp.append(brproperty.getSumProperty(region))
    return tmp
def calculate_direction_color(x,n=3):
    x = np.array(x)/np.linalg.norm(x)
    color =np.array([[1,0,0],[0,1,0],[0,0,1]])
    y = 0
    for i in range(n):
        if x[i]>0:
            y += x[i]*color[i]
        else:
            y += -x[i]*(1-color[i])
    return y/n
def region_layer_seperate(x):
    if x == '':
        return ['','']
    region = ''
    layer = ''
    if x[-1] == '1':
        region = x[:-1]
        layer = '1'
    elif x[-3:] == '2/3':
        region = x[:-3]
        layer = '2/3'
    elif x[-1] == '2':
        region = x[:-1]
        layer = '2'
    elif x[-1] == '3':
        region = x[:-1]
        layer = '3'
    elif x[-1] == '4':
        region = x[:-1]
        layer = '4'
    elif x[-1] == '5':
        region = x[:-1]
        layer = '5'
    elif x[-1] == '6':
        region = x[:-1]
        layer = '6'
    elif x[-2:] in ['6a','6b']:
        region = x[:-2]
        layer = '6'
    else:
        region = x
        layer = ''
    return [region,layer]
def topology_info_compute_by_axon_center(subcerebral_neuron_information_with_py,region_index,topology_region_name):
    tmp = br.getRegion(topology_region_name)
    tmp.annotation()
    region_mask = np.where(tmp.readdata>0,1,0)
    projection_topology_info = {}
    for neuron in subcerebral_neuron_information_with_py:
        if subcerebral_neuron_information_with_py[neuron]['projection_matrix'][region_index]==0:
            continue
        projection_topology_info[neuron] = []
        neuron_tree = iondata.getNeuronTreeByID(neuron[:-7],neuron[-7:])
        for i in neuron_tree.points:
            point = get_mirrored_point(i.xyz)
            if if_inside_mask(region_mask,point/10):
                projection_topology_info[neuron].append(point)
    return projection_topology_info
def topology_info_compute_by_axon(subcerebral_neuron_information_with_py,region_index,topology_region_name):
    tmp = br.getRegion(topology_region_name)
    tmp.annotation()
    region_mask = np.where(tmp.readdata>0,1,0)
    projection_topology_info = {}
    for neuron in subcerebral_neuron_information_with_py:
        if subcerebral_neuron_information_with_py[neuron]['projection_matrix'][region_index]==0:
            continue
        projection_topology_info[neuron] = {'in':[],'out':[]}
        neuron_tree = iondata.getNeuronTreeByID(neuron[:-7],neuron[-7:])
        for edge in neuron_tree.edges:
            edge_mark = -1
            last_point = []
            for i in edge.data:
                point = get_mirrored_point(i.xyz)
                if if_inside_mask(region_mask,point/10):
                    if edge_mark==0:
                        projection_topology_info[neuron]['in'].append(point)
                        edge_mark = 1
                    else:
                        pass
                else:
                    if edge_mark!=0:
                        if edge_mark!=-1:
                            projection_topology_info[neuron]['out'].append(last_point)
                        edge_mark = 0
                    else:
                        pass
                last_point = point
    return projection_topology_info
def get_region_mask(region):
    _br = copy.deepcopy(br)
    tmp = _br.getRegion(region)
    tmp.annotation()
    result = tmp.readdata*1
    del tmp.readdata
    return result
def get_axon_mean_point(neuron_tree,mask,sampling=1):
    points = []
    for point in neuron_tree.points[::sampling]:
        points.append(get_mirrored_point(point.xyz))
    points = np.array(points)/10
    tmp = []
    for i in points:
        if if_inside_mask(mask,i)==1:
            tmp.append(i)
    return [] if len(tmp)==0 else np.mean(tmp,axis=0)
'''Flatten map'''
def compute_depth(x):
    y = []
    if calculating_depth_flatten_position_mode==0:
        tmp_1 = BL.ComputeStreamlines(grid,dv0,dv1,dv2,np.array(x)/20)
        for tmp2 in tmp_1:
            y.append(tmp2[0])
    else:
        for i in x:
            a = np.round(np.array(i)/20).astype(int)
            y.append(depth20[a[0],a[1],a[2]])
    return y
def compute_flatten_position_and_depth(x):
    y = []
    if calculating_depth_flatten_position_mode==0:
        tmp_1 = BL.ComputeStreamlines(grid,dv0,dv1,dv2,np.array(x)/20)
        for tmp2 in tmp_1:
            tmp_3 = Flatmap.map2Flatmap(flatenPara,np.array(tmp2[1])*2,True)
            y.append([tmp2[0],tmp2[1],tmp2[2],tmp_3])
    else:
        for i in x:
            a = np.round(np.array(i)/20).astype(int)
            y.append([depth20[a[0],a[1],a[2]],[outp1[a[0],a[1],a[2]],outp2[a[0],a[1],a[2]],outp3[a[0],a[1],a[2]]],[inp1[a[0],a[1],a[2]],inp2[a[0],a[1],a[2]],inp3[a[0],a[1],a[2]]],[outpx[a[0],a[1],a[2]],outpy[a[0],a[1],a[2]]]])
    return y
'''Visualization sectors'''
def set_view(neuronvis,view):
    if view=='dorsal':
        neuronvis.render.setView('vontral')
    elif view=='left':
        neuronvis.render.setView('right')
    elif view=='right':
        neuronvis.render.setView('left')
    elif view=='ventral':
        neuronvis.render.setView('dorsal')
    else:
        neuronvis.render.setView(view)
def draw_points(visul,points,colors):
    ga = GeometryAdapter.GeometryAdapter()
    ga.geometry.drawModel='points'
    tmp = 0
    for i,j in zip(points,colors):
        ga.geometry.addPoint(i,j)
        ga.geometry.addIndex(tmp)
        tmp+=1
    ga.geometry.type = 1
    visul.render.addGeometry(ga.geometry)
def draw_lines(visul,lines,colors):
    ga = GeometryAdapter.GeometryAdapter()
    ga.geometry.drawModel='lines'
    tmp = -1
    for i,j in zip(lines,colors):
        if len(i)<2:
            continue
        ga.geometry.addPoint(i[0],j[0])
        tmp+=1
        for k,l in zip(i[1:],j[1:]):
            ga.geometry.addPoint(k,l)
            tmp+=1
            ga.geometry.addIndex(tmp-1)
            ga.geometry.addIndex(tmp)
    ga.geometry.type = 1
    visul.render.addGeometry(ga.geometry)
def cluster_visualization(IDs,X,clusters,sampling=0,specific_cluster=-1,colors=[],name='cx'):
    clustering = AgglomerativeClustering(clusters).fit(X)
    if len(colors)==0:
        background_color1 = [0.0,0.0,0.0]
        background_color2 = [1.0,1.0,1.0]
        colors = distinctipy.get_colors(clusters,[background_color1,background_color2])
        distinctipy.color_swatch(colors)
    neuronvis = nv.neuronVis()
    for i,j in zip(IDs,clustering.labels_):
        if specific_cluster!=-1:
            if j!=specific_cluster:
                continue
        if np.random.rand()<sampling:
            continue
        neuronvis.addNeuronByID(i[0],i[1],somaColor=colors[j],axonColor=colors[j],mirrorToRight=True,dendriteHide=True)
    for i in ['vontral','right','anterior']:
        neuronvis.render.setView(i)
        neuronvis.render.savepng(name+'_'+i+'.png')
    neuronvis.render.run()
def cluster_visualization_one_by_one(IDs,X,clusters,sampling=0,specific_cluster=-1,colors=[],name='cx'):
    clustering = AgglomerativeClustering(clusters).fit(X)
    if len(colors)==0:
        background_color1 = [0.0,0.0,0.0]
        background_color2 = [1.0,1.0,1.0]
        colors = distinctipy.get_colors(clusters,[background_color1,background_color2])
        distinctipy.color_swatch(colors)
    neuronvis = nv.neuronVis()
    tmp = {}
    for i,j in zip(IDs,clustering.labels_):
        if specific_cluster!=-1:
            if j!=specific_cluster:
                continue
        if np.random.rand()<sampling:
            continue
        if j not in tmp:
            tmp[j] = []
        tmp[j].append([i[0],i[1]])
    for i in tmp:
        neuronvis.clear()
        for j in tmp[i]:
            neuronvis.addNeuronByID(j[0],j[1],somaColor=colors[i],axonColor=colors[i],mirrorToRight=True,dendriteHide=True)
        for j in ['vontral','left','anterior']:
            neuronvis.render.setView(j)
            neuronvis.render.savepng(name+'_'+str(i)+'_'+j+'.png')
def get_line(neuron_tree,terminal,color=[1,0,0],mark_mirror=False,Z=[]):
    X = []
    Y = []
    edge = neuron_tree.getEdgeByTerminal(terminal)
    while edge != None:
        x = []
        y = []
        for i in edge.data:
            tmp = i.xyz.copy()
            if mark_mirror:
                tmp[2] = 11400-tmp[2]
            x.append(tmp)
            y.append(color)
        if x not in Z:
            X.append(x)
            Y.append(y)
        edge = neuron_tree.getEdgeByTerminal(edge.data[0])
    return [X,Y]
def get_terminal_in_mask(neurons,mask,color=[1,0,0],mode=0):
    X = []
    Y = []
    X_terminal = []
    Y_terminal = []
    for neuron_index,neuron in enumerate(neurons):
        control_console_output(0)
        neuron_tree = get_neuron_tree(neuron)
        control_console_output(1)
        mark_mirror = neuron_tree.root.z>5700
        for terminal in neuron_tree.terminals:
            if terminal.type==3:
                continue
            if if_inside_mask(mask,np.array(terminal.xyz)/10):
                X.append([])
                Y.append([])
                X_terminal.append(terminal.xyz.copy())
                if mark_mirror:
                    X_terminal[-1][2] = 11400-X_terminal[-1][2]
                Y_terminal.append(color if mode==0 else color[neuron_index])
    return [X,Y,X_terminal,Y_terminal]
def set_color(X,color):
    color = np.array(color).tolist()
    Y = []
    for i in X:
        Y.append([])
        for j in i:
            Y[-1].append(color)
    return Y
def project_point_on_line_segment(point, line_start, line_end):
    v = (line_end - line_start).astype(float)
    v /= np.linalg.norm(v)
    w = (point - line_start).astype(float)
    mu = np.dot(w,v)
    projection = line_start + mu*v
    return projection
def get_pdf(mask,neuron_to_brainstem_information,info,r=200,s=125):
    result = None
    kernel = generate_gaussian_kernel(r,s)
    soma_list = None
    for i in info:
        d,p = info[i][:2]
        e = info[i][2:]
        if result==None:
            result = [mask*0.0 for j in range(len(e))]
            soma_list = [[[],[]] for j in range(len(e))]
        for j in range(len(e)):
            fill_kernel(result[j],np.array(p[::-1])*2,e[j]/neuron_to_brainstem_information[i]['relative_density'],r,kernel)
            soma_list[j][0].append(np.array(p[::-1])*2)
            soma_list[j][1].append(e[j])
    for i in range(len(soma_list)):
        soma_list[i][1] = np.array(soma_list[i][1])**0.5
        soma_list[i][1] /= (np.max(soma_list[i][1]) if np.max(soma_list[i][1])!=0 else 1)
    return result,soma_list
def get_terminal_lines_in_mask(neurons,mask,color=[1,0,0],mode=0):
    X = []
    Y = []
    X_terminal = []
    Y_terminal = []
    for neuron_index,neuron in enumerate(neurons):
        control_console_output(0)
        neuron_tree = get_neuron_tree(neuron)
        control_console_output(1)
        mark_mirror = neuron_tree.root.z>5700
        for terminal in neuron_tree.terminals:
            if terminal.type==3:
                continue
            if if_inside_mask(mask,np.array(terminal.xyz)/10):
                X.append([])
                Y.append([])
                a,b = get_line(neuron_tree,terminal,color if mode==0 else color[neuron_index],mark_mirror,X)
                #X.append(a)
                #Y.append(b)
                X += a
                Y += b
                X_terminal.append(terminal.xyz.copy())
                if mark_mirror:
                    X_terminal[-1][2] = 11400-X_terminal[-1][2]
                Y_terminal.append(color if mode==0 else color[neuron_index])
    return [X,Y,X_terminal,Y_terminal]
def get_terminal_lines_in_region(neurons,region_list,color=[1,0,0],mode=0):
    mask = 0.0
    control_console_output(0)
    for i in region_list:
        mask += get_region_mask(i)
    control_console_output(1)
    return get_terminal_lines_in_mask(neurons,mask,color=color,mode=mode)
'''Scene sectors'''
def save_scene(X,color,file_name,region=[],mirror=[],mode=0):
    if mode==0:
        tmp = []
        for i in X:
            tmp.append({'sampleid':i[:-7],'name':i[-7:],'mirror':False,'color':{'r':str(int(color[0]*255)),'g':str(int(color[1]*255)),'b':str(int(color[2]*255))}})
        Scene.createScene(tmp,file_name+'.nv')
    elif mode==1:
        tmp = []
        for i,j in zip(X,color):
            tmp.append({'sampleid':i[:-7],'name':i[-7:],'mirror':False,'color':{'r':str(int(j[0]*255)),'g':str(int(j[1]*255)),'b':str(int(j[2]*255))}})
        Scene.createScene(tmp,file_name+'.nv')
    elif mode==2:
        tmp = []
        for i,j,k in zip(X,color,region):
            tmp.append({'sampleid':i[:-7],'name':i[-7:],'mirror':False,'color':{'r':str(int(j[0]*255)),'g':str(int(j[1]*255)),'b':str(int(j[2]*255))},'soma':k})
        Scene.createScene(tmp,file_name+'.nv')
    elif mode==3:
        tmp = []
        for i,j in zip(X,mirror):
            tmp.append({'sampleid':i[:-7],'name':i[-7:],'mirror':j,'color':{'r':str(int(color[0]*255)),'g':str(int(color[1]*255)),'b':str(int(color[2]*255))}})
        Scene.createScene(tmp,file_name+'.nv')
    else:
        pass
def get_branches(neuron_list,mask,color=[1,0,0],sampling=20):
    result = [[],[]]
    for neuron in tool_bar(neuron_list):
        neuron_tree = iondata.getNeuronTreeByID(neuron[:-7],neuron[-7:])
        for edge in neuron_tree.edges:
            if edge.data[0].type==3:
                continue
            small_edge = []
            samll_edge_color = []
            mark = 0
            for point in edge.data[::sampling]+[edge.data[-1]]:
                position = get_mirrored_point(point.xyz)
                if_inside = if_inside_mask(mask,position/10)
                if mark==0 and if_inside:
                    mark = 1
                    small_edge.append(point.xyz)
                    samll_edge_color.append(color)
                if mark==1 and if_inside:
                    small_edge.append(point.xyz)
                    samll_edge_color.append(color)
                if mark==1 and if_inside==0:
                    if len(small_edge)>1:
                        result[0].append(small_edge)
                        result[1].append(samll_edge_color)
                    small_edge = []
                    samll_edge_color = []
                    mark = 0
                else:
                    pass
    return result
def get_point(p):
    y = np.array(p.xyz)
    return y
def get_terminals(neuron_tree):
    X = [[],[]]
    for terminal in neuron_tree.terminals:
        if terminal.type==3:
            continue
        X[0].append(get_point(terminal))
        X[1].append(terminal)
    return X
def get_branch_points(neuron_tree):
    X = [[],[]]
    for edge in neuron_tree.edges:
        if edge.data[0].type==3:
            continue
        if edge.data[0]==neuron_tree.root:
            continue
        X[0].append(get_point(edge.data[0]))
        X[1].append(edge.data[0])
    return X
def get_line_from_terminal(terminal,color=[1,1,0]):
    X = [[],[]]
    tmp = terminal
    while tmp!=None:
        X[0].append(get_point(tmp))
        X[1].append(color)
        tmp = tmp.parent
    X[0] = X[0][::-1]
    return X
def calculate_angle(a,b):
    if np.sum(a)==0 or np.sum(b)==0:
        return 0
    else:
        return math.acos(np.dot(a,b))*180/np.pi
class cx_edge:
    def __init__(self,x,pre_direction_backward=np.array([0,0,0]),gamma=0.99,angle_threshold=45,parent=None):
        self.points = np.array(x)
        if len(self.points)<2:
            self.direction_forward = np.array([0,0,0])
            self.direction_backward = np.array([0,0,0])
            self.length = 0
        else:
            weights = []
            weight = 1
            for i in range(len(x)-1):
                weights.append(weight)
                weight*=gamma
                if weight<0.09:
                    weight = 0
            weights = np.array(weights).reshape(-1,1)
            directions = self.points[1:]-self.points[0]
            directions /= np.linalg.norm(directions,axis=1).reshape(-1,1)
            direction = np.sum(directions*weights,axis=0)/np.sum(weights)
            direction /= np.linalg.norm(direction)
            self.direction_forward = direction
            directions = self.points[-1]-self.points[:-1][::-1]
            directions /= np.linalg.norm(directions,axis=1).reshape(-1,1)
            direction = np.sum(directions*weights,axis=0)/np.sum(weights)
            direction /= np.linalg.norm(direction)
            self.direction_backward = direction
            self.length = np.sum(np.linalg.norm(self.points[1:]-self.points[:-1],axis=1))
        self.children = []
        self.parent = parent
        self.angle = calculate_angle(pre_direction_backward,self.direction_forward)
        self.is_shaft = True if self.angle<angle_threshold else False
def get_edge(p,start=[],pre_direction_backward=np.array([0,0,0]),angle_threshold=45,parent=None,mark='all'):
    x = start.copy()
    while 1:
        x.append(p.xyz)
        children = []
        for i in p.children:
            if mark=='all':
                children.append(i)
            else:
                if i.index in mark:
                    children.append(i)
        if len(children)==0:
            return cx_edge(x,pre_direction_backward,angle_threshold=angle_threshold,parent=parent)
        elif len(children)==1:
            p = children[0]
        else:
            tmp = cx_edge(x,pre_direction_backward,angle_threshold=angle_threshold,parent=parent)
            for i in children:
                tmp.children.append(get_edge(i,[p.xyz],tmp.direction_backward,angle_threshold=angle_threshold,parent=tmp,mark=mark))
            tmp2 = []
            for i in tmp.children:
                tmp2.append(i.length)
            tmp.length += np.max(tmp2)
            return tmp
def get_shaft(neuron_tree,shaft_color=[1,0,0],nonshaft_color=[0,1,0],mode=1,return_nonshaft=False,angle_threshold=45,length_ratio_threshold=0.8,length_threshold=-1,mirror=False,mark='all'):
    '''
        mode 0 : maximum length path
        mode 1 : large pathes
        mode 2 : mode 1 + angle check
    '''
    mirror_mark = (neuron_tree.root.z>5700 and mirror)
    _edge = get_edge(neuron_tree.root,angle_threshold=angle_threshold,mark=mark)
    lines = []
    lines_color = []
    lines2 = []
    lines_color2 = []
    if mode==0:
        Q = [_edge]
        shaft = []
        while Q!=[]:
            edge = Q.pop(0)
            shaft.append(edge)
            tmp = []
            for p in edge.points:
                tmp.append(shaft_color)
            lines.append(abs(np.array([0,0,11400])-edge.points) if mirror_mark else edge.points)
            lines_color.append(tmp)
            _length = []
            for tmp in edge.children:
                _length.append(tmp.length)
            if _length!=[]:
                Q.append(edge.children[np.argmax(_length)])
        if return_nonshaft:
            Q = [_edge]
            while Q!=[]:
                edge = Q.pop(0)
                if edge not in shaft:
                    tmp = []
                    for p in edge.points:
                        tmp.append(nonshaft_color)
                    lines2.append(abs(np.array([0,0,11400])-edge.points) if mirror_mark else edge.points)
                    lines_color2.append(tmp)
                for tmp in edge.children:
                    Q.append(tmp)
    elif mode<=2:
        Q = [_edge]
        shaft = []
        while Q!=[]:
            edge = Q.pop(0)
            shaft.append(edge)
            tmp = []
            for p in edge.points:
                tmp.append(shaft_color)
            lines.append(abs(np.array([0,0,11400])-edge.points) if mirror_mark else edge.points)
            lines_color.append(tmp)
            _max_length = -1
            for tmp in edge.children:
                _max_length = np.max([_max_length,tmp.length])
            if _max_length>=length_threshold:
                for tmp in edge.children:
                    if tmp.length/_max_length>=length_ratio_threshold:
                        if mode==1:
                            Q.append(tmp)
                        else:
                            if tmp.is_shaft:
                                Q.append(tmp)
        if return_nonshaft:
            Q = [_edge]
            while Q!=[]:
                edge = Q.pop(0)
                if edge not in shaft:
                    tmp = []
                    for p in edge.points:
                        tmp.append(nonshaft_color)
                    lines2.append(abs(np.array([0,0,11400])-edge.points) if mirror_mark else edge.points)
                    lines_color2.append(tmp)
                for tmp in edge.children:
                    Q.append(tmp)
    else:
        pass
    return [[lines,lines_color],[lines2,lines_color2]]
def get_points_from_point(x):
    tmp = []
    while x!=None:
        tmp.append(x.index)
        x = x.parent
    return tmp
